#!pip install faiss-gpu
import faiss
import torch
import glob,os,numpy as np,pickle
from PIL import Image
from torchvision import transforms,models
import matplotlib.pyplot as plt


val_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

faiss_index = faiss.IndexFlatL2(1000)  # build the index
model = models.resnet18().cuda()
path='checkpoints/Hg20220810141015/9.845261573791504_278.pth'
model.load_state_dict(torch.load(path))
print('load sucess!')

# storing the image representations
im_indices = []

PATH_TRAIN = "/home/hegang/datas2/hegang/datas/public_datasets/contrast_learning/trainData"
PATH_TEST = "/home/hegang/datas2/hegang/datas/public_datasets/contrast_learning/testData"

with torch.no_grad():
    for f in glob.glob(os.path.join(PATH_TRAIN, '*/*')):
        im = Image.open(f)
        im = im.resize((224, 224))
        im = torch.tensor([val_transforms(im).numpy()]).cuda()

        preds = model(im)
        preds = np.array([preds[0].cpu().numpy()])
        faiss_index.add(preds)  # add the representation to index
        im_indices.append(f)  # store the image name to find it later on

f_Index=open('IndexFlatL2.pkl','wb')
pickle.dump(faiss_index, f_Index, protocol = 4)
#Retrieval with a query image
with torch.no_grad():
    for f in os.listdir(PATH_TEST):
        # query/test image
        print(os.path.join(PATH_TEST, f))
        im = Image.open(os.path.join(PATH_TEST, f))
        im_=im.copy()
        im = im.resize((224, 224))
        im = torch.tensor([val_transforms(im).numpy()]).cuda()
        test_embed = model(im).cpu().numpy()
        print(test_embed.shape)
        _, I = faiss_index.search(test_embed, 1000)
        print("Retrieved Image: {}".format(im_indices[I[0][0]]))

        search_res=Image.open(im_indices[I[0][0]])


        plt.figure()
        plt.subplot(1,2,1)
        plt.imshow(im_)
        plt.subplot(1, 2, 2)
        plt.imshow(search_res)
        plt.axis('off')  # 不显示坐标轴
        plt.show()
        plt.close()
